from sqlalchemy import (
    create_engine, Column, Integer, String, Enum, TIMESTAMP, ForeignKey, Text, JSON, BigInteger, SmallInteger
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship, sessionmaker
from sqlalchemy.sql import text
import bcrypt
import os

# 修改为你的数据库连接信息
DATABASE_URL = "mysql+pymysql://root:zixi1004@localhost:3306/rbac_system?charset=utf8mb4"

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    username = Column(String(50), unique=True, nullable=False)
    email = Column(String(100), unique=True, nullable=False)
    password_hash = Column(String(255), nullable=False)
    salt = Column(String(64), nullable=False)
    user_type = Column(Enum('student', 'teacher', 'staff', 'admin'), default='staff', nullable=False)
    department = Column(String(50), default='general')
    status = Column(Enum('active', 'disabled'), default='active')
    created_at = Column(TIMESTAMP)

class UserAttribute(Base):
    __tablename__ = 'user_attributes'
    user_id = Column(Integer, ForeignKey('users.id', ondelete='CASCADE'), primary_key=True)
    attr_key = Column(String(50), primary_key=True)
    attr_value = Column(String(255))

class Role(Base):
    __tablename__ = 'roles'
    id = Column(Integer, primary_key=True)
    role_name = Column(String(50), unique=True, nullable=False)
    description = Column(String(255))

class UserRole(Base):
    __tablename__ = 'user_roles'
    user_id = Column(Integer, ForeignKey('users.id', ondelete='CASCADE'), primary_key=True)
    role_id = Column(Integer, ForeignKey('roles.id', ondelete='CASCADE'), primary_key=True)

class Permission(Base):
    __tablename__ = 'permissions'
    id = Column(Integer, primary_key=True)
    resource = Column(String(50), nullable=False)
    action = Column(String(50), nullable=False)
    attribute_condition = Column(JSON)
    __table_args__ = (
        # 联合唯一约束
        {'mysql_engine': 'InnoDB'},
    )

class RolePermission(Base):
    __tablename__ = 'role_permissions'
    role_id = Column(Integer, ForeignKey('roles.id', ondelete='CASCADE'), primary_key=True)
    permission_id = Column(Integer, ForeignKey('permissions.id', ondelete='CASCADE'), primary_key=True)

class OperationLog(Base):
    __tablename__ = 'operation_logs'
    id = Column(BigInteger, primary_key=True, autoincrement=True)
    user_id = Column(Integer, ForeignKey('users.id', ondelete='SET NULL'), nullable=True)
    session_id = Column(String(100))
    operation_type = Column(String(50), nullable=False)
    operation_target = Column(String(100))
    operation_time = Column(TIMESTAMP)
    operation_details = Column(Text)
    ip_address = Column(String(45))
    user_agent = Column(String(255))
    status = Column(Enum('success', 'fail'), default='success')

class LogOperationType(Base):
    __tablename__ = 'log_operation_types'
    id = Column(SmallInteger, primary_key=True, autoincrement=True)
    type_code = Column(String(50), unique=True, nullable=False)
    description = Column(String(100), nullable=False)

class LogDetail(Base):
    __tablename__ = 'log_details'
    log_id = Column(BigInteger, ForeignKey('operation_logs.id', ondelete='CASCADE'), primary_key=True)
    detail_key = Column(String(50), primary_key=True)
    detail_value = Column(Text)

def hash_password(plain_password: str) -> tuple:
    salt = bcrypt.gensalt()
    hashed = bcrypt.hashpw(plain_password.encode('utf-8'), salt)
    return hashed.decode('utf-8'), salt.decode('utf-8')

# 创建数据库连接和表
def main():
    engine = create_engine(DATABASE_URL, echo=True)
    # 创建所有表
    Base.metadata.create_all(engine)

    # 检查并自动添加 attribute_condition 字段
    with engine.connect() as conn:
        result = conn.execute(text(
            "SHOW COLUMNS FROM permissions LIKE 'attribute_condition';"
        ))
        if result.fetchone() is None:
            conn.execute(text(
                "ALTER TABLE permissions ADD COLUMN attribute_condition JSON NULL COMMENT '属性条件表达式，支持ABAC';"
            ))

    # 检查并自动添加 status 字段
    with engine.connect() as conn:
        result = conn.execute(text(
            "SHOW COLUMNS FROM users LIKE 'status';"
        ))
        if result.fetchone() is None:
            conn.execute(text(
                "ALTER TABLE users ADD COLUMN status ENUM('active', 'disabled') DEFAULT 'active' COMMENT '用户状态';"
            ))

    # 创建视图和插入初始数据
    with engine.connect() as conn:
        # 插入常用操作类型
        conn.execute(text("""
        INSERT IGNORE INTO log_operation_types (type_code, description) VALUES
        ('login', '用户登录'),
        ('logout', '用户登出'),
        ('access_denied', '权限拒绝'),
        ('resource_access', '资源访问'),
        ('data_update', '数据更新'),
        ('data_delete', '数据删除'),
        ('pwd_change', '密码修改'),
        ('role_assign', '角色分配');
        """))
        # 插入用户（使用单向加盐加密），并获取id
        user_ids = {}
        users = [
            ('alice', 'alice@example.com', 'alice123', 'student', 'math', 'active'),
            ('bob', 'bob@example.com', 'bob456', 'teacher', 'physics', 'active'),
            ('carol', 'carol@example.com', 'carol789', 'admin', 'it', 'active'),
        ]
        for username, email, plain_pwd, user_type, department, status in users:
            pwd_hash, salt = hash_password(plain_pwd)
            result = conn.execute(text("""
                INSERT INTO users (username, email, password_hash, salt, user_type, department, status)
                VALUES (:username, :email, :password_hash, :salt, :user_type, :department, :status)
                ON DUPLICATE KEY UPDATE id=LAST_INSERT_ID(id)
            """), {
                'username': username,
                'email': email,
                'password_hash': pwd_hash,
                'salt': salt,
                'user_type': user_type,
                'department': department,
                'status': status
            })
            user_id = conn.execute(text("SELECT id FROM users WHERE username=:username"), {'username': username}).scalar()
            user_ids[username] = user_id

        # 插入角色，并获取id
        role_ids = {}
        roles = [
            ('student', '学生'),
            ('teacher', '教师'),
            ('admin', '管理员'),
        ]
        for role_name, description in roles:
            conn.execute(text("""
                INSERT INTO roles (role_name, description)
                VALUES (:role_name, :description)
                ON DUPLICATE KEY UPDATE id=LAST_INSERT_ID(id)
            """), {'role_name': role_name, 'description': description})
            role_id = conn.execute(text("SELECT id FROM roles WHERE role_name=:role_name"), {'role_name': role_name}).scalar()
            role_ids[role_name] = role_id

        # 插入用户-角色关联
        conn.execute(text("""
            INSERT IGNORE INTO user_roles (user_id, role_id) VALUES
            (:alice, :student), (:bob, :teacher), (:carol, :admin)
        """), {
            'alice': user_ids['alice'],
            'student': role_ids['student'],
            'bob': user_ids['bob'],
            'teacher': role_ids['teacher'],
            'carol': user_ids['carol'],
            'admin': role_ids['admin'],
        })

        # 插入权限，并获取id
        permission_ids = {}
        permissions = [
            ('user_management', 'user_management', '{"user_type": "admin"}'),
            ('role_management', 'role_management', '{"role_name": "admin"}'),
            ('permission_management', 'permission_management', '{"resource": "admin"}'),
        ]
        for resource, action, attribute_condition in permissions:
            conn.execute(text("""
                INSERT INTO permissions (resource, action, attribute_condition)
                VALUES (:resource, :action, :attribute_condition)
                ON DUPLICATE KEY UPDATE id=LAST_INSERT_ID(id)
            """), {
                'resource': resource,
                'action': action,
                'attribute_condition': attribute_condition
            })
            permission_id = conn.execute(text("SELECT id FROM permissions WHERE resource=:resource AND action=:action"), {'resource': resource, 'action': action}).scalar()
            permission_ids[f"{resource}_{action}"] = permission_id

        # 插入角色-权限关联
        conn.execute(text("""
            INSERT IGNORE INTO role_permissions (role_id, permission_id) VALUES
            (:student, :user_management), (:teacher, :role_management), (:admin, :permission_management)
        """), {
            'student': role_ids['student'],
            'user_management': permission_ids['user_management_user_management'],
            'teacher': role_ids['teacher'],
            'role_management': permission_ids['role_management_role_management'],
            'admin': role_ids['admin'],
            'permission_management': permission_ids['permission_management_permission_management'],
        })

        # 插入用户属性
        conn.execute(text("""
            INSERT IGNORE INTO user_attributes (user_id, attr_key, attr_value) VALUES
            (:alice, 'department', 'math'), (:bob, 'department', 'physics'), (:carol, 'department', 'it')
        """), {
            'alice': user_ids['alice'],
            'bob': user_ids['bob'],
            'carol': user_ids['carol'],
        })

        # 创建视图
        conn.execute(text("""
        CREATE OR REPLACE VIEW v_user_permissions AS
        SELECT
            u.id AS user_id,
            u.username,
            r.role_name,
            p.resource,
            p.action,
            p.id AS permission_id,
            p.attribute_condition
        FROM users u
        JOIN user_roles ur ON u.id = ur.user_id
        JOIN roles r ON ur.role_id = r.id
        JOIN role_permissions rp ON r.id = rp.role_id
        JOIN permissions p ON rp.permission_id = p.id;
        """))
    print("数据库结构和视图已创建完毕！")

if __name__ == "__main__":
    main() 